import numpy as np
import matplotlib.pyplot as plt
import tqdm

def sourlas_encode(source_code: np.array, code_rate, partial_a):
    N = len(source_code)
    varia_node={} # i to a
    func_node={} # a to i

    # generate random factor graph
    for i in range(round(N/code_rate)):
        b = np.random.choice(range(N), partial_a, replace=False)
        func_node[i]=b
        for j in b:
            if j not in varia_node:
                varia_node[j]=[]
            varia_node[j].append(i)
    
    # encode on basis of factor graph
    codeword={}
    for funcnode in func_node.keys():
        j=1
        for varianode in func_node[funcnode]:
            j *= source_code[varianode]
        codeword[funcnode]=j
    return varia_node, func_node, codeword

def bsc_channel(codeword, p_error):
    received = codeword.copy()
    n_flips = round(len(codeword)*p_error)
    flip_idx = np.random.randint(0, len(codeword), size=n_flips)
    for i in flip_idx:
        received[i] = 1 - received[i]
    return received

def sourlas_decode(received, beta, partial_a, varia_node, func_node, num_bits):
    def iteration(v,f,J,beta,N,max_iter=500,tol=1e-4,K=partial_a):
        flag=False
        h={}
        u={}
        #np.random.seed(42)
        for varianode in v.keys():#h_{i\to a}初始化
            for funcnode in v[varianode]:
                if varianode not in h:
                    h[varianode] = {}
                #h[varianode][funcnode]=np.random.normal(0,1,1)
                h[varianode][funcnode] = np.random.uniform(-1, 1, 1)
                #h[varianode][funcnode] = np.random.choice((-1,1), 1)
        for funcnode in f.keys():#u_{a\to i}初始化
            for varianode in f[funcnode]:
                if funcnode not in u:
                    u[funcnode] = {}
                #u[funcnode][varianode]=np.random.normal(0,1,1)
                u[funcnode][varianode] = np.random.uniform(-1, 1, 1)
                #u[funcnode][varianode] = np.random.choice((-1,1), 1)
        for t in range(max_iter):#迭代更新
            delta=0.0
            for i in v.keys():
                for a in v[i]:#h_{ia}
                    h_new=0
                    for b in v[i]:#连接i的所有节点
                        if b!=a:
                            h_new += u[b][i]
                            #print(u[b][i])
                    delta+=abs(h_new-h[i][a])
                    h[i][a]=h_new

            for a in f.keys():
                for i in f[a]:#u_{ai}
                    tanhprod=1
                    for j in f[a]:#连接a的所有节点
                        if j!=i:
                            tanhprod *= np.tanh(beta*h[j][a])
                    u_new=(np.arctanh(np.tanh(beta*J[a])*tanhprod))/(beta)
                    delta+=abs(u_new-u[a][i])
                    u[a][i]=u_new


            if delta/(N*K)<tol:
                flag=True
                break
        # if flag:
            # print(t,'time Converged')
        # else:
            # print('Not Converged')
        return h,u,flag
    
    def decode(v,u,beta,N):
        decoded_code=np.zeros(N)
        #print(f'u is\n',u,'\n')
        for i in range(N):
            sum_val=0
            if i not in v:
                continue
            for b in v[i]:
                #print(u[funcnode][i])
                sum_val+=u[b][i]
            # print(sum_val)
            decoded_code[i]=np.sign(sum_val)

        #print(f'decoded code is\n',decoded_code)
        return decoded_code
    h, u, flag = iteration(varia_node,func_node,received,beta,num_bits)
    decoded = decode(varia_node,u,beta,num_bits)
    return decoded, flag

def calc_accuracy(source_code, decoded_code):
    N = len(source_code)
    accu = np.sum(source_code == decoded_code)/N
    return accu

def perform_test_once():
    source_code=np.random.choice([-1,1], num_bits)
    varia_node, func_node, codeword = sourlas_encode(source_code, code_rate, partial_a)
    # print(varia_node, '\n\n\n', func_node)
    received = bsc_channel(codeword, p_error)
    decoded, flag = sourlas_decode(received, beta, partial_a, varia_node, func_node, num_bits)

    print(f"信息位: {source_code}")
    print(f"编码: {np.array([received[i] for i in sorted(received.keys())])}")
    print(f"接收: {np.array([received[i] for i in sorted(received.keys())])}")
    print(f"解码: {decoded}")

def plt_accuracy():
    def main(p_error,beta):
        source_code=np.random.choice([-1,1], num_bits)
        varia_node, func_node, codeword = sourlas_encode(source_code, code_rate, partial_a)
        received = bsc_channel(codeword, p_error)
        decoded, flag = sourlas_decode(received, beta, partial_a, varia_node, func_node, num_bits)
        o = calc_accuracy(source_code, decoded)
        return o,flag
    
    P=np.arange(0.01,0.4,0.01)
    accuracy=[]
    for p in tqdm.tqdm(P):
        beta=np.log((1-p)/p)*0.5
        current_accuracy = 0
        for _ in range(10):
            o,flag=main(p,beta)
            if flag:
                if current_accuracy == 0:
                    current_accuracy = o
                current_accuracy = (current_accuracy+o)/2
        accuracy.append(np.mean(current_accuracy))

    plt.figure()
    plt.xlabel('flipping rate p')
    plt.ylabel('accuracy')
    plt.plot(P,accuracy)
    plt.title(r'$\beta_c=\frac{1}{2}\ln{\frac{1-p}{p}}$')
    plt.show()


if __name__ == "__main__":
    # ----- args -----
    num_bits = 50
    partial_a = 3
    code_rate = 0.5
    p_error = 0.03 # possibility of bit flip
    beta = 0.5*np.log((1-p_error)/p_error)
    # -----------------
    # perform_test_once()
    plt_accuracy()

